from .models import get_model
from .database import FaissDB, SqliteDB


class SearchEngine(object):
    def __init__(self, config):
        self.config = config
        self.face_feature = get_model(config.MODEL)
        self.faiss_db = FaissDB(config.FAISS)
        self.sqlite_db = SqliteDB(config.SQLITE)
        self.match_thresh = config.MATCH_THRESHOLD
    
    def add_face(self, image_path, person_name):
        """
        assert added face has single face with known identity
        """
        if self.sqlite_db.exists(image_path):
            # skip this image if it is already added.
            return
        
        # extract feature
        feat = self.face_feature.run(image_path)
        if feat is None:
            print("failed to parse image {}".format(image_path))
            return
        if len(feat) != 1:
            print('{} has {} faces, not added'.format(image_path, len(feat)))
            return
        feat = feat[0]
        
        # insert into database
        faiss_id = self.faiss_db.insert(feat)
        self.sqlite_db.insert(faiss_id, image_path, person_name)
    
    def exists(self, image_path):
        return self.sqlite_db.exists(image_path)
    
    def search_by_image(self, image_path, topk=1):
        features = self.face_feature.run(image_path)

        if features is None:
            return {'num_faces': 0, "data": []}

        num_faces = 0
        all_face_data = []
        for feat in features:
            # search
            distances, indices = self.faiss_db.search(feat, topk)
            # match result of the first input.
            distances = distances[0]
            indices = indices[0]

            face_match = []
            for count, index in enumerate(indices):
                data = self.sqlite_db.search(index)
                if data is None: continue
                match_path = data[0]
                match_name = data[1]
                match_conf = float(distances[count])
                if match_conf < self.match_thresh:
                    continue
                # face_match.append({'image_path': match_path, 'name': match_name, 'confidence': match_conf})
                # return without facedb image path
                face_match.append({'name': match_name, 'confidence': match_conf})
            
            all_face_data.append(face_match)
            num_faces += 1
        
        all_faces = {'num_faces': num_faces, "data": all_face_data}

        return all_faces
    
    def save(self):
        self.faiss_db.save_index(self.config.FAISS.INDEX_PATH)
        # self.sqlite_db.close()

    def count(self):
        n1 = self.faiss_db.count()
        n2 = self.sqlite_db.count()
        assert (n1 == n2)
        return n1
